"""
使用说明:
告诉 Hugging Face Datasets 库在离线模式下工作。
当设置 HF_DATASETS_OFFLINE=1 后，该库将不会尝试在线下载数据集，而是会尝试在本地找到数据集的缓存版本。
export HF_DATASETS_OFFLINE=1
python start_evaluation.py
"""

print("----------------------------------------------加载数据集--------------------------------------------------")
import os
import sys
# 加载数据集
from datasets import load_from_disk

# 从环境变量中获取数据集路径
dataset_path = os.getenv('DATASET_PATH')

if not dataset_path:
    print("未找到数据集路径，请设置环境变量 DATASET_PATH 并指定数据集的路径")
    sys.exit(1)  # 退出程序，并返回状态码 1 表示错误

# 加载数据集
dataset = load_from_disk(dataset_path)

# 修改数据集中的列名
dataset_dict = dataset.rename_columns({"groundTruths": "ground_truth"})

# 划分数据集
dataset_dict = dataset_dict.train_test_split(test_size=0.5)  # 80% 训练集，20% 验证集
print(f"划分数据集\n{dataset_dict}")

print("-----------------------------------------------加载模型---------------------------------------------------")

from langchain_openai import ChatOpenAI, OpenAIEmbeddings

BASE_URL = "http://172.16.176.59:7002"

# 加载大语言模型
llm = ChatOpenAI(base_url=BASE_URL, api_key="Empty", model="Baichuan2-13B-Chat")
print("加载大语言模型完成")

embedding_model = OpenAIEmbeddings(base_url=BASE_URL, api_key="Empty", model="bge-large-zh-v1.5", tiktoken_enabled=False, tiktoken_model_name="E:/LLM/tiktoken/tiktoken_cl100k_base")
print("加载 embedding 模型完成")

print("----------------------------------------------评估数据集--------------------------------------------------")

from ragas import evaluate
from ragas.metrics import (
    answer_relevancy,
    answer_similarity,
    context_entity_recall,
    context_precision,
    faithfulness,
)

# 使用ragas库的evaluate函数评估模型
result = evaluate(
    dataset=dataset_dict["train"],  # 传入训练集
    llm=llm,
    raise_exceptions=False,
    embeddings=embedding_model,  # 传入 embedding
    metrics=[
        answer_relevancy,
        answer_similarity,
        context_precision,
        context_entity_recall,
        faithfulness,
    ],
)
print(f"""
评估结果 / Evaluation Results:\n
回答相关性：     {result["answer_relevancy"]}
回答相似度：     {result["answer_similarity"]}
上下文精确度：   {result["context_precision"]}
上下文实体召回率：{result["context_entity_recall"]}
忠实度：        {result["faithfulness"]}
""")
